﻿// Copyright (c) Microsoft. All rights reserved.

using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Linq;
using Microsoft.SemanticKernel.Connectors.HuggingFace.Core;

namespace Microsoft.SemanticKernel.Connectors.HuggingFace;

/// <summary>
/// Represents the metadata of a Hugging Face chat completion.
/// </summary>
public sealed class HuggingFaceTextGenerationStreamMetadata : ReadOnlyDictionary<string, object?>
{
    internal HuggingFaceTextGenerationStreamMetadata() : base(new Dictionary<string, object?>()) { }

    internal HuggingFaceTextGenerationStreamMetadata(TextGenerationStreamResponse streamResponse) : this()
    {
        this.Index = streamResponse.Index;
        this.TokenId = streamResponse.Token?.Id ?? 0;
        this.TokenSpecial = streamResponse.Token?.Special;
        this.TokenLogProb = streamResponse.Token?.LogProb;
        this.GeneratedText = streamResponse.GeneratedText;
        this.GeneratedTokens = streamResponse.Details?.GeneratedTokens;
        this.FinishReason = streamResponse.Details?.FinishReason;
    }

    private HuggingFaceTextGenerationStreamMetadata(IDictionary<string, object?> dictionary) : base(dictionary) { }

    /// <summary>
    /// Index of the chunk
    /// </summary>
    public int Index
    {
        get => this.GetValueFromDictionary(nameof(this.Index)) as int? ?? 0;
        internal init => this.SetValueInDictionary(value, nameof(this.Index));
    }

    /// <summary>
    /// Token identifier.
    /// </summary>
    public int TokenId
    {
        get => this.GetValueFromDictionary(nameof(this.TokenId)) as int? ?? 0;
        internal init => this.SetValueInDictionary(value, nameof(this.TokenId));
    }

    /// <summary>
    /// Special flag
    /// </summary>
    public bool? TokenSpecial
    {
        get => this.GetValueFromDictionary(nameof(this.TokenSpecial)) as bool? ?? false;
        internal init => this.SetValueInDictionary(value, nameof(this.TokenSpecial));
    }

    /// <summary>
    /// The log probabilities of the completion.
    /// </summary>
    public double? TokenLogProb
    {
        get => this.GetValueFromDictionary(nameof(this.TokenLogProb)) as double? ?? 0;
        internal init => this.SetValueInDictionary(value, nameof(this.TokenLogProb));
    }

    /// <summary>
    /// Text generated by the model.
    /// </summary>
    public string? GeneratedText
    {
        get => this.GetValueFromDictionary(nameof(this.GeneratedText)) as string;
        internal init => this.SetValueInDictionary(value, nameof(this.GeneratedText));
    }

    /// <summary>
    /// Number of generated tokens.
    /// </summary>
    public int? GeneratedTokens
    {
        get => this.GetValueFromDictionary(nameof(this.GeneratedTokens)) as int?;
        internal init => this.SetValueInDictionary(value, nameof(this.GeneratedTokens));
    }

    /// <summary>
    /// Finish reason.
    /// </summary>
    public string? FinishReason
    {
        get => this.GetValueFromDictionary(nameof(this.FinishReason)) as string;
        internal init => this.SetValueInDictionary(value, nameof(this.FinishReason));
    }

    /// <summary>
    /// Converts a dictionary to a <see cref="HuggingFaceChatCompletionMetadata"/> object.
    /// </summary>
    public static HuggingFaceTextGenerationStreamMetadata FromDictionary(IReadOnlyDictionary<string, object?> dictionary) => dictionary switch
    {
        null => throw new ArgumentNullException(nameof(dictionary)),
        HuggingFaceTextGenerationStreamMetadata metadata => metadata,
        IDictionary<string, object?> metadata => new HuggingFaceTextGenerationStreamMetadata(metadata),
        _ => new HuggingFaceTextGenerationStreamMetadata(dictionary.ToDictionary(pair => pair.Key, pair => pair.Value))
    };

    private void SetValueInDictionary(object? value, string propertyName)
        => this.Dictionary[propertyName] = value;

    private object? GetValueFromDictionary(string propertyName)
        => this.Dictionary.TryGetValue(propertyName, out var value) ? value : null;
}
